#pragma once

#include "vkinfer/runtime/operator.h"
#include "vkinfer/runtime/vk/vkblob.h"
#include "vkinfer/runtime/vk/vkcontext.h"

namespace vkinfer
{
    class VkOp : public Operator
    {
    protected:
        VkDevice device;
        std::shared_ptr<VkCompute> cmd;
        std::shared_ptr<VkTransfer> transfer;

        VkDescriptorSetLayout descriptorSetLayout;
        VkDescriptorPool descriptorPool;
        VkDescriptorSet descriptorSet;

        VkShaderModule shaderModule;
        VkPipelineLayout pipelineLayout;
        // VkPipelineCache pipelineCache;
        VkPipeline pipeline;

    public:
        ~VkOp();
        void initialize(VkDevice device, const std::string& shader_file, uint32_t num_outputs, uint32_t num_inputs,
            uint32_t num_weights, uint32_t pushconstants_size);
        void set_context(const std::shared_ptr<VkCompute>& cmd, const std::shared_ptr<VkTransfer>& transfer);
        
        // update io buffers
        virtual void update_descriptors(
            const std::vector<std::shared_ptr<VkBlob>>& inputs,
            std::vector<std::shared_ptr<VkBlob>>& outputs) = 0;

    private:
        void load_shader(const std::string& shaderfile);
        void create_descriptor_set_layout(uint32_t num_descriptors);
        void create_descriptor_set(uint32_t num_descriptors);
        void create_pipeline(uint32_t pushconstants_size);
    };

}
